import operator
from typing import Annotated, TypedDict, List

from langchain_core.messages import BaseMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END

# --- 员工代理 (子图) 的定义 ---
# 这是一个简单的代理，接收问题并回答


class EmployeeState(TypedDict):
    """子图的状态"""

    question: str
    answer: str


def employee_node(state: EmployeeState) -> dict:
    """员工代理的核心逻辑：接收问题，生成答案"""
    question = state["question"]
    # 在真实应用中，这里会调用 LLM 来生成回答
    # 为简化教程，我们只返回一个格式化的字符串
    answer = f"I have answered your question: '{question}'"
    print(f"--- EMPLOYEE: answered question: '{question}' ---")
    return {"answer": answer}


# 创建子图的构建器
employee_builder = StateGraph(EmployeeState)
# 添加唯一的节点
employee_builder.add_node("responder", employee_node)
# 定义子图的入口和出口
employee_builder.set_entry_point("responder")
employee_builder.set_finish_point("responder")

# 编译子图，现在它是一个可调用的 Runnable 对象
# 我们可以像调用普通函数一样调用它：employee_graph.invoke({"question": "..."})
employee_graph = employee_builder.compile()


# --- 主管代理 (父图) 的定义 ---
# 这个代理负责将一系列问题分发给员工代理


class SupervisorState(TypedDict):
    """父图的状态"""

    questions: List[str]
    # 使用 Annotated 和 operator.add 来累积答案
    answers: Annotated[List[str], operator.add]
    # 追踪当前处理到哪个问题
    current_question_index: int


def supervisor_entry_node(state: SupervisorState) -> dict:
    """主管代理的入口节点，做一些初始化工作"""
    print("--- SUPERVISOR: Starting the process. ---")
    # 初始化问题索引
    return {"current_question_index": 0}


def call_employee_node(state: SupervisorState) -> dict:
    """
    这是包装节点，是连接父图和子图的桥梁。
    它负责调用员工代理（子图）。
    """
    # 1. 从父图状态中获取当前需要处理的问题
    current_index = state["current_question_index"]
    current_question = state["questions"][current_index]
    print(f"--- SUPERVISOR: Delegating question '{current_question}' to employee. ---")

    # 2. 准备子图需要的输入状态
    child_state_input = {"question": current_question}

    # 3. 调用子图！
    child_state_output = employee_graph.invoke(child_state_input)

    # 4. 从子图的输出中提取结果
    answer = child_state_output["answer"]

    # 5. 将结果包装成父图状态的更新
    return {"answers": [answer], "current_question_index": current_index + 1}


def supervisor_router(state: SupervisorState) -> str:
    """
    主管代理的路由节点。
    它决定是继续处理下一个问题，还是结束流程。
    """
    print("--- SUPERVISOR: Making a decision. ---")
    num_questions = len(state["questions"])
    current_index = state["current_question_index"]

    if current_index >= num_questions:
        print("--- SUPERVISOR: All questions answered. Finishing. ---")
        return "end"
    else:
        print(f"--- SUPERVISOR: More questions to go. Continuing. ---")
        return "continue"


# 创建父图的构建器
supervisor_builder = StateGraph(SupervisorState)

# 添加父图的节点
supervisor_builder.add_node("supervisor_start", supervisor_entry_node)
supervisor_builder.add_node("call_employee", call_employee_node)

# 设置父图的入口点
supervisor_builder.set_entry_point("supervisor_start")

# 添加从入口到调用员工节点的边
supervisor_builder.add_edge("supervisor_start", "call_employee")

# 添加条件边，实现循环
supervisor_builder.add_conditional_edges(
    # 起始节点是 "call_employee"
    "call_employee",
    # 路由函数是 supervisor_router
    supervisor_router,
    # 映射关系
    {
        "continue": "call_employee",  # 如果返回 "continue"，则再次调用员工节点
        "end": END,  # 如果返回 "end"，则结束图的执行
    },
)

# 编译父图
supervisor_graph = supervisor_builder.compile()


if __name__ == "__main__":
    # 定义我们要处理的一系列问题
    questions_to_ask = [
        "What is the weather today?",
        "What is the capital of France?",
        "How does LangGraph work?",
    ]

    # 定义父图的初始输入
    # 注意，我们只需要提供 questions，其他字段会自动初始化或在流程中生成
    initial_input = {"questions": questions_to_ask}

    # 调用父图来执行整个流程
    final_state = supervisor_graph.invoke(initial_input)

    print("\n--- Final Result ---")
    print(f"Initial Questions: {final_state['questions']}")
    print(f"Final Answers: {final_state['answers']}")
    print(f"Process finished after {final_state['current_question_index']} questions.")

    # 你也可以可视化这个图，看到它的循环结构
    # from IPython.display import Image, display
    # try:
    #     img_data = supervisor_graph.get_graph().draw_mermaid_png()
    #     display(Image(img_data))
    # except Exception:
    #     print("Failed to draw graph. Please install pydot and graphviz.") 